/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * Global Trust Authority is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use crate::entities::db_model::rv_db_model::Model as RvDbModel;
use crate::entities::inner_model::rv_model::{RefValueModel};
use crate::entities::request_body::rv_del_req_body::RvDelReqBody;
use crate::entities::request_body::rv_update_req_body::RvUpdateReqBody;
use crate::error::ref_value_error::RefValueError;
use crate::error::ref_value_error::RefValueError::{DbError, InvalidParameter};
use crate::repositories::rv_db_repo::RvDbRepo;
use crate::repositories::rv_dtl_db_repo::RvDtlDbRepo;
use crate::services::rv_trait::RefValueTrait;
use crate::utils::utils::Utils;
use actix_web::web::{Data};
use key_management::api::{CryptoOperations, DefaultCryptoImpl};
use log::{error};
use rdb::get_connection;
use sea_orm::{DatabaseConnection, TransactionTrait};
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use config_manager::types::CONFIG;

pub struct RvMysqlImpl;

impl RvMysqlImpl {
    pub(crate) fn new() -> Self {
        Self
    }
}

impl RefValueTrait for RvMysqlImpl {
    async fn add(&self, conn: Data<Arc<DatabaseConnection>>, rv_model: &RefValueModel) -> Result<(), RefValueError> {
        let raw_conn: &DatabaseConnection = &**conn;
        let limit = CONFIG.get_instance().unwrap().attestation_service.ref_value.single_user_ref_value_limit as u64;
        match RvDbRepo::add(raw_conn, rv_model, limit).await {
            Ok(_) => {
                let txn = conn.begin().await.map_err(|e| DbError(e.to_string()))?;
                RvDtlDbRepo::add(&txn, rv_model).await?;
                txn.commit().await.map_err(|e| DbError(e.to_string()))?;
                Ok(())
            }
            Err(e) => {
                error!("Ref value add failed: {}", e);
                Err(e)
            }
        }
    }

    async fn update(
        &self,
        conn: Data<Arc<DatabaseConnection>>,
        user_id: &str,
        update_req_body: &RvUpdateReqBody,
    ) -> Result<RvDbModel, RefValueError> {
        let txn = conn.begin().await.map_err(|e| DbError(e.to_string()))?;
        let updated_model = RvDbRepo::update(&txn, &user_id, &update_req_body).await.map_err(|e| {
            error!("Reference value update failed: {}", e);
            e
        })?;

        let new_content = update_req_body
            .content
            .as_ref()
            .map(|content| content.to_string())
            .unwrap_or(String::new());
        let new_type = update_req_body
            .attester_type
            .as_ref()
            .map(|attester_type| attester_type.to_string())
            .unwrap_or(String::new());
        if !new_content.is_empty() {
            let mut details = Utils::parse_rv_detail_from_jwt_content(&new_content)?;
            details.set_all_ids(&update_req_body.id);
            details.set_uid(user_id);
            details.set_attester_type(&updated_model.attester_type);
            RvDtlDbRepo::del_by_rv_ids(&txn, user_id, &vec![update_req_body.id.clone()]).await?;
            let chunks = details.reference_values.chunks(500);
            for chunk in chunks {
                RvDtlDbRepo::add_dtls(&txn, chunk.into()).await?;
            }
        } else if !new_type.is_empty() {
            RvDtlDbRepo::update_type_by_rv_id(&txn, user_id, &update_req_body.id, &new_type).await?;
        }
        txn.commit().await.map_err(|e| RefValueError::DbError(e.to_string()))?;
        Ok(updated_model)
    }

    async fn delete(&self, conn: Data<Arc<DatabaseConnection>>, user_id: &str, del_type: &str, del_req_body: &RvDelReqBody) -> Result<(), RefValueError> {
        let del_req_body_clone = del_req_body.clone();
        match del_type {
            "all" => RvDbRepo::del_all(&conn, &user_id).await,
            "id" => {
                let txn = conn.begin().await.map_err(|e| RefValueError::DbError(e.to_string()))?;
                RvDbRepo::del_by_id(&txn, &user_id, &del_req_body_clone.ids.clone().unwrap()).await?;
                RvDtlDbRepo::del_by_rv_ids(&txn, &user_id, &del_req_body_clone.ids.unwrap()).await?;
                txn.commit().await.map_err(|e| RefValueError::DbError(e.to_string()))?;
                Ok(())
            }
            "type" => {
                let txn = conn.begin().await.map_err(|e| RefValueError::DbError(e.to_string()))?;
                RvDbRepo::del_by_type(&conn, &user_id, &del_req_body_clone.attester_type.clone().unwrap()).await?;
                RvDtlDbRepo::del_by_attester_type(&txn, &user_id, &del_req_body_clone.attester_type.unwrap()).await?;
                txn.commit().await.map_err(|e| RefValueError::DbError(e.to_string()))?;
                Ok(())
            }
            _ => Err(InvalidParameter(format!("Invalid delete_type: {}", del_req_body.delete_type))),
        }
    }

    fn verify<'a>(&'a self,
                  measurements: &'a Vec<String>,
                  user_id: &'a str,
                  attester_type: &'a str
    ) -> Pin<Box<dyn Future<Output = Result<Vec<String>, String>> + Send + 'a>> {
        Box::pin(async move {
            let conn = match get_connection().await {
                Ok(conn) => conn,
                Err(e) => {
                    let error_msg = format!("Database connection failed: {}", e);
                    error!("{}", error_msg);
                    return Err(error_msg);
                }
            };

            // Clone the measurements to own them for the async block
            let measurements_ref: Vec<&str> = measurements.iter().map(|s| s.as_str()).collect();

            let is_require_sign = match Utils::is_require_sign() {
                Ok(val) => val,
                Err(e) => return Err(e.message()),
            };
            // Call the appropriate verification method based on configuration
            let result = if is_require_sign {
                // Verify with signature validation
                Self::verify_with_sign(&conn, measurements_ref, user_id, attester_type).await
            } else {
                // Verify without signature validation
                Self::verify_without_sign(&conn, measurements_ref, user_id, attester_type).await
            };

            Ok(result)
        })
    }
}

impl RvMysqlImpl {
    /// Verify signature for a single reference value
    ///
    /// # Arguments
    /// * `model` - The complete reference value model
    ///
    /// # Returns
    /// * `bool` - true if signature is valid, false otherwise
    async fn verify_single_rv_signature(model: &crate::entities::db_model::rv_db_model::Model) -> bool {
        use sea_orm::IntoActiveModel;

        let active_model = model.clone().into_active_model();

        let data = match Utils::encode_rv_db_model_to_bytes(active_model) {
            Ok(data) => data,
            Err(e) => {
                error!("Failed to encode reference value for signature verification: {}", e);
                return false;
            }
        };

        match DefaultCryptoImpl
            .verify("FSK", Some(&model.key_version), data, model.signature.clone())
            .await
        {
            Ok(true) => true,
            Ok(false) => {
                error!("Signature verification failed for rv_id: {}", model.id);
                false
            }
            Err(e) => {
                error!("Failed to verify reference value signature: {}", e);
                false
            }
        }
    }

    /// Extract SHA256 hashes from JWT content
    ///
    /// # Arguments
    /// * `content` - The reference value JWT content
    ///
    /// # Returns
    /// * `Result<HashSet<String>, RefValueError>` - Set of SHA256 hashes found in the content
    fn extract_sha256_hashes_from_content(content: &str) -> Result<HashSet<String>, RefValueError> {
        let details = Utils::parse_rv_detail_from_jwt_content(content)?;

        let hash_set: HashSet<String> = details.reference_values
            .into_iter()
            .map(|rv_dtl| rv_dtl.sha256)
            .collect();

        Ok(hash_set)
    }

    pub async fn verify_without_sign(
        conn: &DatabaseConnection,
        measurements: Vec<&str>,
        user_id: &str,
        attester_type: &str,
    ) -> Vec<String> {
        let chunk_set: Vec<String> = measurements.iter().map(|s| ToString::to_string(&s)).collect();
        match RvDtlDbRepo::query_missing_sha256(conn, chunk_set, attester_type, user_id).await {
            Ok(unmatched) => {unmatched}
            Err(e) => {
                error!("Verification failed, msg: {}", e.to_string());
                Vec::new()
            }
        }
    }

    pub async fn verify_with_sign(
        conn: &DatabaseConnection,
        measurements: Vec<&str>,
        user_id: &str,
        attester_type: &str,
    ) -> Vec<String> {
        let measurements_strings: Vec<String> = measurements.iter().map(|s| s.to_string()).collect();

        // Step 1: Batch query all rv_ids for the measurements
        let hash_to_rv_id = match RvDtlDbRepo::query_rv_ids_by_measurements(
            conn, &measurements_strings, user_id, attester_type
        ).await {
            Ok(mapping) => mapping,
            Err(e) => {
                error!("Failed to query rv_ids by measurements: {}", e);
                return measurements.into_iter().map(String::from).collect();
            }
        };

        if hash_to_rv_id.is_empty() {
            return measurements.into_iter().map(String::from).collect();
        }

        // Step 2: Get unique rv_ids and batch query their complete models
        let unique_rv_ids: Vec<String> = hash_to_rv_id.values().cloned().collect::<std::collections::HashSet<_>>().into_iter().collect();

        let rv_models: HashMap<String, RvDbModel> = match RvDbRepo::query_content_and_signature_by_ids(conn, &unique_rv_ids).await {
            Ok(data) => data,
            Err(e) => {
                error!("Failed to query content and signatures: {}", e);
                return measurements.into_iter().map(String::from).collect();
            }
        };

        // Step 3: Verify signatures and extract valid hash sets
        let mut verified_hash_sets: HashMap<String, HashSet<String>> = HashMap::new();

        for (rv_id, model) in rv_models {
            // Verify signature using the complete model
            if Self::verify_single_rv_signature(&model).await {
                // Extract hash set from signed content
                match Self::extract_sha256_hashes_from_content(&model.content) {
                    Ok(hash_set) => {
                        verified_hash_sets.insert(rv_id, hash_set);
                    }
                    Err(e) => {
                        error!("Failed to extract hashes from content for rv_id {}: {}", rv_id, e);
                    }
                }
            } else {
                error!("Signature verification failed for rv_id: {}", rv_id);
            }
        }

        // Step 4: Check which measurements are valid
        let mut matched = HashSet::new();

        for measurement in &measurements_strings {
            if let Some(rv_id) = hash_to_rv_id.get(measurement) {
                if let Some(hash_set) = verified_hash_sets.get(rv_id) {
                    if hash_set.contains(measurement) {
                        matched.insert(measurement.clone());
                    }
                }
            }
        }

        // Step 5: Return unmatched measurements
        measurements_strings.into_iter()
            .filter(|measurement| !matched.contains(measurement))
            .collect()
    }
}
